import torchvision
from torch.utils.data import DataLoader

from constants import *
from utils.file_util import *
from data_process.data_loader import *
from utils.critirion import *
import pandas as pd

import torch.nn.functional as F

from model.ultra_attn_model import *
from utils.network_util import *
from model.base_model import *
from android_model.ultra_attn_model import AndroidUltraNetwork


class LipReading():
    def __init__(self, paras):
        self.corpusPath = paras['corpusPath']
        self.charPath = paras['charPath']
        self.trainPath = paras['trainPath']
        self.testPath = paras['testPath']
        self.contentPath = paras['contentPath']

        self.trainBatch = paras['trainBatch']
        self.testBatch = paras['testBatch']
        self.device = paras['device']

        self.mode = paras['mode']

        self.learnRate = paras['learnRate']
        self.clip = paras['clip']
        self.epoch = paras['epoch']
        self.saveEvery = paras['saveEvery']
        self.checkPath = paras['checkPath']
        self.testWer = 10000.0
        self.testLoss = 10000.0

        self.netType = paras['netType']
        self.vocabType = paras['vocabType']

        self.modelType = paras['modelType']
        self.modelFunc = paras['modelFunc']
        self.curVolun = paras['curVolun']
        self.trainNumber = paras['trainNumber']
        self.dataMode = paras['dataMode']


        print(self.device)

        self.change = False

        # self.charPath=charPath

        # 字符索引映射表
        self.s2lDict, self.l2sDict = CreateMap(self.charPath)
        self.classNum = len(self.s2lDict)

        self.matrixDir = 'confuse_matrix'
        if not os.path.exists(self.matrixDir):
            os.makedirs(self.matrixDir)

        # 定义网络模型
        self.initModel(self.modelType)
        # self.model = TestConv(self.classNum)
        self.model = self.model.to(self.device)

        #所有loss
        self.CELossFn = nn.CrossEntropyLoss()
        self.KDLossFn = kdLossFunction
        self.MSELossFn = nn.MSELoss()

        # self.opt = optim.Adadelta(self.model.parameters())
        self.opt = torch.optim.SGD(self.model.parameters(), lr=0.1,
                                   momentum=0.9, weight_decay=5e-4)
        # self.opt = torch.optim.Adam(self.model.parameters())

        # 获得根目录路径
        self.lastPath = ''

    def initModel(self, modelType):
        if modelType=='resnet1':
            self.model = UltraNetwork(self.classNum, firstChannel=1)
            # self.model = AndroidUltraNetwork(self.classNum, firstChannel=1)
        elif modelType == 'resnet8':
            self.model = UltraNetwork(self.classNum, firstChannel=8)
        elif modelType=='alexnet':
            self.model = AlexNetUltra(self.classNum)
        elif modelType=='vgg':
            self.model= VGGNet(self.classNum)
        elif modelType == 'mobilenet':
            self.model = MobileNetv2(self.classNum)
        else:
            raise ValueError

    # def PrintHint(self):
    #     if self.modelType=='resnet':
    #         print('current type is resnet')
    #     if self.modelType=='vgg':
    #         print('current type is vgg')


    def GetLoader(self, pathFile, batch):
        if 'train' in pathFile and self.dataMode == DataModoType.SINGLE:
            dataset = MultiDataset(pathFile)
        else:
            dataset = MyDataset(pathFile)
        loader = DataLoader(dataset, batch_size=batch, shuffle=True, num_workers=0)
        return loader

    def trainDealData(self, loader, curEpoch):
        allLoss = 0
        allWer = 0
        count = 0
        confuseMatrix = np.zeros((self.classNum, self.classNum))
        for idx, sample in enumerate(loader):

            audioInput = sample['audioInput']
            labelInput = sample['labelInput']
            # 将数据放到device上
            audioInput = audioInput.to(self.device)
            labelInput = labelInput.to(self.device)
            # 需要先将数据更改成目标数据的shape
            audioInput = audioInput.float()
            labelInput = labelInput.long()

            if self.modelType in NORMAL_MODEL:
                output, loss = self.normalModelDealData(audioInput, labelInput)
            elif self.modelType in ATTENTION_MODEL:
                output, loss = self.attentionModelDealData(audioInput, labelInput)
            else:
                raise ValueError

            curCount = len(labelInput)
            count += curCount

            self.opt.zero_grad()
            loss.backward()
            # _ = nn.utils.clip_grad_norm_(self.model.parameters(), self.clip)
            self.opt.step()
            allLoss += loss * curCount
            targetSent, predSent = self.getTranslation(output, labelInput, idx == len(loader) - 1)
            curWer = EvalWer(targetSent, predSent)
            allWer += curWer

            confuseMatrix = updateConfuseMatrix(confuseMatrix, predSent, targetSent, self.s2lDict)

        print(f'train sample is {count}')
        printMsg = 'epoch {}, loss is {:.4f}, wer is {:.4f}\n\n'.format(curEpoch, allLoss / count, allWer / count)
        print(printMsg)

        lossPath = 'train_loss.txt'
        with open(lossPath, 'a+') as f:
            msg = 'epoch {}, loss is {:.4f}, wer is {:.4f}\n'.format(curEpoch, allLoss / count, allWer / count)
            f.write(msg)

        self.confuseMatrix = confuseMatrix

    def testDealData(self, loader, lossFn, onlyTest=False):
        allLoss = 0
        allWer = 0
        count = 0

        if onlyTest:
            saveFile = open('test_res.txt', 'w')

        for idx, sample in enumerate(loader):

            audioInput = sample['audioInput']
            labelInput = sample['labelInput']
            # 将数据放到device上
            audioInput = audioInput.to(self.device)
            labelInput = labelInput.to(self.device)
            # 需要先将数据更改成目标数据的shape
            audioInput = audioInput.float()
            labelInput = labelInput.long()

            if self.modelType in NORMAL_MODEL:
                output, loss = self.normalModelDealData(audioInput, labelInput)
            elif self.modelType in ATTENTION_MODEL:
                output, loss = self.attentionModelDealData(audioInput, labelInput)
            else:
                raise ValueError

            curCount = len(labelInput)
            count += curCount

            loss = lossFn(output, labelInput)

            allLoss += loss * curCount
            targetSent, predSent = self.getTranslation(output, labelInput, idx == len(loader) - 1)
            if onlyTest:
                saveFile.write(f'{targetSent} {predSent}\n')
            curWer = EvalWer(targetSent, predSent)
            allWer += curWer

        testWer = allWer / count
        testLoss = allLoss / count
        print(f'test sample is {count}')
        printMsg = 'test loss is {:.4f}, wer is {:.4f}\n\n'.format(allLoss / count, allWer / count)
        print(printMsg)
        if not onlyTest:
            self.saveToFile(testWer, testLoss)

    def normalModelDealData(self, audioInput, labelInput):
        output = self.model(audioInput)
        ce_loss = self.CELossFn(output, labelInput)
        loss = ce_loss
        return output, loss

    def attentionModelDealData(self, audioInput, labelInput):
        output, middle_output1, middle_output2, middle_output3, final_fea, \
        middle1_fea, middle2_fea, middle3_fea, attn1, attn2, attn3, attn4 = self.model(audioInput)

        temperature = 3
        alpha = 0.1
        beta = 1e-6
        ce_loss = self.CELossFn(output, labelInput)
        if self.modelFunc == 'kl':
            # print('come to kl')
            temp4 = output / temperature
            temp4 = torch.softmax(temp4, dim=1)
            loss1by4 = self.KDLossFn(middle_output1, temp4.detach(), temperature) * (temperature ** 2)
            loss2by4 = self.KDLossFn(middle_output2, temp4.detach(), temperature) * (temperature ** 2)
            loss3by4 = self.KDLossFn(middle_output3, temp4.detach(), temperature) * (temperature ** 2)
            loss = (1 - alpha) * ce_loss + \
                   alpha * (loss1by4 + loss2by4 + loss3by4)
        elif self.modelFunc == 'attn':
            # print('come to attn')
            batch = attn1.shape[0]
            attn1Shape = attn1.shape[-2:]
            attn2Shape = attn2.shape[-2:]
            attn3Shape = attn3.shape[-2:]
            attn4Shape = attn4.shape[-2:]
            changeAttn1 = F.adaptive_avg_pool2d(attn1, attn2Shape)
            changeAttn2 = F.adaptive_avg_pool2d(attn2, attn3Shape)
            changeAttn3 = F.adaptive_avg_pool2d(attn3, attn4Shape)
            attn_loss1 = self.MSELossFn(changeAttn1, attn2.detach())
            attn_loss2 = self.MSELossFn(changeAttn2, attn3.detach())
            attn_loss3 = self.MSELossFn(changeAttn3, attn4.detach())

            loss = (1 - alpha) * ce_loss + \
                   beta * (attn_loss1 + attn_loss2 + attn_loss3)
        elif self.modelFunc == 'kl_attn':
            # print('come to kl and attn')
            temp4 = output / temperature
            temp4 = torch.softmax(temp4, dim=1)
            loss1by4 = self.KDLossFn(middle_output1, temp4.detach(), temperature) * (temperature ** 2)
            loss2by4 = self.KDLossFn(middle_output2, temp4.detach(), temperature) * (temperature ** 2)
            loss3by4 = self.KDLossFn(middle_output3, temp4.detach(), temperature) * (temperature ** 2)

            batch = attn1.shape[0]
            attn1Shape = attn1.shape[-2:]
            attn2Shape = attn2.shape[-2:]
            attn3Shape = attn3.shape[-2:]
            attn4Shape = attn4.shape[-2:]
            changeAttn1 = F.adaptive_avg_pool2d(attn1, attn2Shape)
            changeAttn2 = F.adaptive_avg_pool2d(attn2, attn3Shape)
            changeAttn3 = F.adaptive_avg_pool2d(attn3, attn4Shape)
            attn_loss1 = self.MSELossFn(changeAttn1, attn2.detach())
            attn_loss2 = self.MSELossFn(changeAttn2, attn3.detach())
            attn_loss3 = self.MSELossFn(changeAttn3, attn4.detach())

            loss = (1 - alpha) * ce_loss + \
                   alpha * (loss1by4 + loss2by4 + loss3by4) + \
                   beta * (attn_loss1 + attn_loss2 + attn_loss3)
        else:
            loss = ce_loss
        return output, loss

    def getTranslation(self, output, labelInput, isPrint):
        predLabel = torch.max(output, dim=-1)[-1]
        predSent = Label2Sent(predLabel, self.l2sDict, self.netType, vocabType=self.vocabType)
        targetSent = Label2Sent(labelInput, self.l2sDict, self.netType, vocabType=self.vocabType)

        if isPrint:
            sentLen = len(targetSent)
            for i in range(sentLen):
                print(f'prediction: {predSent[i]} || target: {targetSent[i]}')
        return targetSent, predSent

    def saveToFile(self, testWer, testLoss):
        if self.testWer > testWer:
            self.testWer = testWer
            self.testLoss = testLoss
            self.change = True

            matrixSavePath = os.path.join(self.matrixDir, f'{self.curVolun}_{self.modelType}-{self.modelFunc}.csv')
            saveTruct = pd.DataFrame(data=self.confuseMatrix)
            saveTruct.to_csv(matrixSavePath, index=False, header=False)

            self.SaveCheck(self.testWer)

    def Train(self, teachRate=1):
        self.startEpoch = 0
        self.model.train()
        torch.cuda.empty_cache()

        if self.checkPath!=None:
            self.LoadCheck(self.checkPath)

        trainLoader = self.GetLoader(self.trainPath, self.trainBatch)

        for t in range(self.startEpoch, self.epoch):
            self.change = False
            print(f'current epoch is {t}')
            print('train process')

            self.adjustLearningRate(t)
            self.trainDealData(trainLoader, t)
            # 顺便做测试
            with torch.no_grad():
                self.Test()

            if self.change:
                with open('testWer.txt', 'a+') as f:
                    f.write(f'test wer is {self.testWer}, cur epoch is {t}\n')
            # 判断是否收敛，收敛就退出
        vocabPath = os.path.splitext(self.charPath)[0].split('/')[-1]
        writeMsg = f'model:{self.modelType}-{self.modelFunc} | vocabPath:{vocabPath} | curVolun:{self.curVolun} | bestWer:{format(self.testWer, ".4f")} | trainNumber:{format(self.trainNumber, ".4f")}\n'
        with open('all_result.txt', 'a+') as f:
            f.write(writeMsg)

    def Test(self, onlyTest=False):
        self.model.train()

        lossFn = nn.CrossEntropyLoss()
        testLoader = self.GetLoader(self.testPath, self.testBatch)
        print('test process')
        # 取出数据，训练
        self.testDealData(testLoader, lossFn, onlyTest)


    def SaveCheck(self, testWer):
        directory = f'ModelWeights'
        if not os.path.exists(directory):
            os.makedirs(directory)
        # savePath = os.path.join(directory, 'loss{:.4f}_wer{:.4f}.pt'.format(testLoss, testWer))
        savePath = os.path.join(directory, f'{self.modelType}-{self.modelFunc}_' + f'{self.curVolun}_' + 'bestWer{:.4f}.pt'.format(testWer))
        torch.save({
            'model': self.model.state_dict(),
            'opt': self.opt.state_dict(),
            'testWer': self.testWer
        }, savePath)

        epochPath = 'ModelWeights/last_check.pt'
        torch.save({
            'model': self.model.state_dict(),
            'opt': self.opt.state_dict(),
            'testWer': self.testWer
        }, epochPath)

        if self.lastPath != '':
            os.remove(self.lastPath)
        self.lastPath = savePath

    def LoadCheck(self, checkPath=None):
        if checkPath != None:
            checkpoint = torch.load(checkPath, map_location=self.device)
            self.model.load_state_dict(checkpoint['model'])
            # self.opt.load_state_dict(checkpoint['opt'])
            # self.testWer = checkpoint['testWer']
            # self.lastPath = checkPath
            print(f'{checkPath} is loaded')

    def adjustLearningRate(self, epoch):
        if False and (epoch < 1):
            lr = 0.01
        # elif 75 <= epoch < 130:
        #     lr = self.args.lr * (self.args.stepRatio ** 1)
        # elif 130 <= epoch < 180:
        #     lr = self.args.lr * (self.args.stepRatio ** 2)
        # elif epoch >= 180:
        #     lr = self.args.lr * (self.args.stepRatio ** 3)
        # else:
        #     lr = self.args.lr
        elif 75 <= epoch < 130:
            lr = 0.1 * (0.1 ** 1)
        elif 130 <= epoch < 180:
            lr = 0.1 * (0.1 ** 2)
        elif epoch >= 180:
            lr = 0.1 * (0.1 ** 3)
        else:
            lr = 0.1

        for param_group in self.opt.param_groups:
            param_group['lr'] = lr


def kdLossFunction(output, target_output,temperature):
    """Compute kd loss"""
    """
    para: output: middle ouptput logits.
    para: target_output: final output has divided by temperature and softmax.
    """

    output = output / temperature
    output_log_softmax = torch.log_softmax(output, dim=1)
    loss_kd = -torch.mean(torch.sum(output_log_softmax * target_output, dim=1))
    return loss_kd